#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Copyright © 启明星辰 版权所有
# @Time    : 2024/6/6 14:10
# @Author  : 胡浩浩
# @File    : copy_file_to_docker.py
# @IDE     : PyCharm
# @description :
import os

import jmespath
import ruamel.yaml
import docker

client = docker.from_env()

def read_yml(file_path, node=None):
    with open(file_path, "r", encoding="utf-8") as f:
        yaml = ruamel.yaml.YAML(typ='unsafe', pure=True)
        yaml.allow_duplicate_keys = True
        data = yaml.load(f)
    if node:
        return data[node]
    else:
        return data if data else []

def copy_file_to_docker(container_id, source_path, target_path):
    try:
        container = client.containers.get(container_id)
        container.start()
        # 重启容器
        # 构建 docker cp 命令
        docker_cp_command = f"docker cp {source_path} {container_id}:{target_path}"
        # 使用 os.system 执行命令
        os.system(docker_cp_command)
        print(f"File '{source_path}' copied to container '{container_id}' successfully.")
        container.start()
        print("File copied successfully and container restarted.")
        return True
    except docker.errors.NotFound as e:
        print(f"Container {container_id} not found: {e}")
        return False
    except docker.errors.APIError as e:
        print(f"Error stopping or removing container {container_id}: {e}")
        return False

if __name__ == '__main__':
    docker_list = read_yml("data_containers_list.yml")
    container_id_list = jmespath.search("@[*][0]", docker_list)
    for container_id in container_id_list:
        copy_file_to_docker(container_id, "source/app.py", "/usr/local/lib/python3.11/site-packages/mitmproxy/tools/web/")
